import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import networkx as nx
from typing import Dict, List, Any, Optional, Tuple
import logging
from app.schemas.facenet_base_schema import TaskStatus
from app.schemas.dhgam.hgcn_schema import HGCNParams, HGCNResult, NodeEmbedding, PredictionResult

# 配置日志
logger = logging.getLogger(__name__)

class HGCNService:
    def __init__(self):
        pass

    async def process_hgcn(self, input_params: Dict[str, Any], task_id: str) -> Dict[str, Any]:
        """
        处理HGCN算法请求
        
        Args:
            input_params: HGCN算法输入参数（字典形式）
            task_id: 任务ID
            
        Returns:
            Dict: 包含任务结果的字典
        """
        
        try:
            # 将字典转换为HGCNParams对象
            params = HGCNParams(**input_params)
            
            # 构建NetworkX图
            G = self._build_networkx_graph(params.graph)
            
            # 检查图是否为空或存在孤立节点
            if len(G.nodes()) == 0:
                raise ValueError("图中没有节点")
            
            # 移除孤立节点，避免零度问题
            isolated_nodes = list(nx.isolates(G))
            if isolated_nodes:
                logger.warning(f"图中存在 {len(isolated_nodes)} 个孤立节点，这些节点将被移除")
                G.remove_nodes_from(isolated_nodes)
                
                # 如果移除后图为空，则抛出异常
                if len(G.nodes()) == 0:
                    raise ValueError("移除孤立节点后，图中没有剩余节点")
            
            # 初始化HGCN模型
            hgcn = HeterogeneousGCN(
                hidden_dims=params.hidden_dims,
                learning_rate=params.learning_rate,
                epochs=params.epochs
            )
            
            # 准备标签（如果有）
            labels = None
            if params.labels:
                # 将字典转换为有序列表，确保顺序与图中的节点一致
                node_list = list(G.nodes())
                labels_list = []
                for node_id in node_list:
                    if node_id in params.labels:
                        labels_list.append(params.labels[node_id])
                    else:
                        # 对于没有标签的节点，使用-1或其他适当的值
                        labels_list.append(-1)
                labels = labels_list
            
            # 训练模型
            hgcn.fit(G, labels=labels, task=params.task)
            
            # 获取节点嵌入
            embeddings = hgcn.get_embeddings()
            
            # 执行预测
            predictions = None
            if params.task == "node_classification":
                try:
                    predictions = hgcn.predict()
                except Exception as predict_error:
                    logger.error(f"节点分类预测失败: {str(predict_error)}", exc_info=True)
                    # 预测失败但不影响整体流程，继续执行
            
            # 为社团检测任务添加自定义逻辑
            elif params.task == "community_detection":
                try:
                    # 将嵌入转换为numpy数组用于聚类
                    embeddings_array = np.array([emb for emb in embeddings.values()])
                    
                    # 使用K-means进行社团检测
                    from sklearn.cluster import KMeans
                    
                    # 估计社团数量（这里使用标签中的不同值作为社团数量）
                    if params.labels:
                        n_clusters = len(set(params.labels.values()))
                    else:
                        # 如果没有标签，使用节点类型数量作为社团数量
                        node_types = set()
                        for node in G.nodes(data=True):
                            if 'type' in node[1]:
                                node_types.add(node[1]['type'])
                        n_clusters = len(node_types)
                    
                    # 确保至少有2个社团
                    n_clusters = max(n_clusters, 2)
                    
                    # 执行K-means聚类
                    kmeans = KMeans(n_clusters=n_clusters, random_state=0)
                    cluster_labels = kmeans.fit_predict(embeddings_array)
                    
                    # 将聚类结果映射回节点ID
                    predictions = {}
                    node_ids = list(embeddings.keys())
                    for i, node_id in enumerate(node_ids):
                        predictions[node_id] = int(cluster_labels[i])
                    
                    logger.info(f"社团检测完成，共检测到 {n_clusters} 个社团")
                
                except Exception as cluster_error:
                    logger.error(f"社团检测失败: {str(cluster_error)}", exc_info=True)
                    # 社团检测失败但不影响整体流程，继续执行
            
            # 格式化结果
            result = self._format_results(embeddings, G, predictions)
            
            # 构建响应
            response = {
                "task_id": task_id,
                "task_status": TaskStatus.COMPLETED,
                "task_progress": 100.0,
                "input_params": input_params,  # 保留原始输入参数
                "output_params": result.dict(),
                "metrics": []
            }
            
            return response
        except Exception as e:
            logger.error(f"HGCN处理失败: {str(e)}", exc_info=True)
            return {
                "task_id": task_id,
                "task_status": TaskStatus.FAILED,
                "task_progress": 0,
                "input_params": input_params,  # 保留原始输入参数
                "error_message": str(e)
            }
    
    def _build_networkx_graph(self, graph_data):
        """
        将请求中的图数据转换为NetworkX图
        
        Args:
            graph_data: 请求中的图结构
            
        Returns:
            nx.Graph: NetworkX图对象
        """
        G = nx.Graph()
        
        # 添加节点
        for node in graph_data.nodes:
            # 添加节点属性
            node_attrs = {"type": node.type.value}
            node_attrs.update(node.properties)
            G.add_node(node.id, **node_attrs)
        
        # 添加边
        for edge in graph_data.edges:
            # 添加边属性
            edge_attrs = {"type": edge.type.value}
            edge_attrs.update(edge.properties)
            G.add_edge(edge.source, edge.target, **edge_attrs)
        
        return G
    
    def _format_results(self, embeddings, graph, predictions=None) -> HGCNResult:
        """
        格式化HGCN结果
        
        Args:
            embeddings: 节点嵌入
            graph: NetworkX图对象
            predictions: 节点分类预测结果 (可选)
            
        Returns:
            HGCNResult: 格式化后的结果
        """
        result_embeddings = []
        prediction_results = []
        
        for node_id, embedding in embeddings.items():
            node_embedding = NodeEmbedding(
                node_id=node_id,
                embedding=embedding.tolist() if isinstance(embedding, np.ndarray) else embedding
            )
            result_embeddings.append(node_embedding)
            
            # 如果有预测结果，添加到输出中
            if predictions and node_id in predictions:
                prediction = PredictionResult(
                    node_id=node_id,
                    predicted_label=int(predictions[node_id]) if isinstance(predictions[node_id], (int, float, np.integer)) else predictions[node_id]
                )
                prediction_results.append(prediction)
        
        return HGCNResult(embeddings=result_embeddings, predictions=prediction_results)

# 从现有的异构图卷积神经网络类导入
from app.services.dhgam.hgcn_service import HeterogeneousGCN 